import os
import json
import random
import shutil


def polygon_to_yolo_bbox(polygon_points, image_width, image_height):
    """
    将多边形 polygon 格式的边界框（bbox）转换为 YOLO 格式的边界框。

    :param polygon_points: 多边形的点坐标列表，每个点是一个包含 x 和 y 坐标的列表，例如 [[x1, y1], [x2, y2], ...]
    :param image_width: 图像的宽度
    :param image_height: 图像的高度
    :return: YOLO 格式的边界框，包含中心点的 x 坐标、中心点的 y 坐标、宽度和高度
    """
    x_coords = [point[0] for point in polygon_points]
    y_coords = [point[1] for point in polygon_points]
    x_min = min(x_coords)
    x_max = max(x_coords)
    y_min = min(y_coords)
    y_max = max(y_coords)
    x_center = (x_min + x_max) / (2 * image_width)
    y_center = (y_min + y_max) / (2 * image_height)
    width = (x_max - x_min) / image_width
    height = (y_max - y_min) / image_height
    return x_center, y_center, width, height


def json_to_yolo(json_file_path, output_dir):
    """
    将 JSON 格式的标注文件转换为 YOLO 格式的标注文件。

    :param json_file_path: JSON 标注文件的路径
    :param output_dir: 输出 YOLO 标注文件的目录
    :return: 生成的 YOLO 标注文件的路径
    """
    with open(json_file_path, 'r') as f:
        data = json.load(f)
    image_width = data['imageWidth']
    image_height = data['imageHeight']
    yolo_lines = []
    for shape in data['shapes']:
        class_id = 0
        polygon_points = shape['points']
        x_center, y_center, width, height = polygon_to_yolo_bbox(polygon_points, image_width, image_height)
        yolo_line = f"{class_id} {x_center} {y_center} {width} {height}"
        yolo_lines.append(yolo_line)
    txt_file_name = os.path.basename(json_file_path).replace('.json', '.txt')
    txt_file_path = os.path.join(output_dir, txt_file_name)
    with open(txt_file_path, 'w') as f:
        f.write('\n'.join(yolo_lines))
    return txt_file_path


def split_dataset(json_dir, image_dir, output_dir, train_ratio=0.8):
    """
    将数据集划分为训练集和验证集，并将对应的图像和标注文件复制到相应的目录中。

    :param json_dir: 包含 JSON 标注文件的目录
    :param image_dir: 包含图像文件的目录
    :param output_dir: 输出划分后数据集的根目录
    :param train_ratio: 训练集所占的比例，默认为 0.8
    """
    # 创建输出目录
    train_image_dir = os.path.join(output_dir, 'images', 'train')
    train_label_dir = os.path.join(output_dir, 'labels', 'train')
    val_image_dir = os.path.join(output_dir, 'images', 'val')
    val_label_dir = os.path.join(output_dir, 'labels', 'val')
    os.makedirs(train_image_dir, exist_ok=True)
    os.makedirs(train_label_dir, exist_ok=True)
    os.makedirs(val_image_dir, exist_ok=True)
    os.makedirs(val_label_dir, exist_ok=True)

    json_files = [f for f in os.listdir(json_dir) if f.endswith('.json')]
    random.shuffle(json_files)  # 随机打乱 JSON 文件列表的顺序
    train_size = int(len(json_files) * train_ratio)  # 计算训练集的大小
    train_files = json_files[:train_size]  # 获取训练集的 JSON 文件列表
    val_files = json_files[train_size:]  # 获取验证集的 JSON 文件列表

    # 处理训练集的文件
    for json_file in train_files:
        json_file_path = os.path.join(json_dir, json_file)
        txt_file_path = json_to_yolo(json_file_path, train_label_dir)  # 将 JSON 标注文件转换为 YOLO 格式
        image_file_name = json_file.replace('.json', '.jpg')  # 生成对应的图像文件名
        image_file_path = os.path.join(image_dir, image_file_name)  # 构建图像文件的完整路径
        shutil.copy(image_file_path, train_image_dir)  # 将图像文件复制到训练集图像目录

    # 处理验证集的文件
    for json_file in val_files:
        json_file_path = os.path.join(json_dir, json_file)
        txt_file_path = json_to_yolo(json_file_path, val_label_dir)  # 将 JSON 标注文件转换为 YOLO 格式
        image_file_name = json_file.replace('.json', '.jpg')  # 生成对应的图像文件名
        image_file_path = os.path.join(image_dir, image_file_name)  # 构建图像文件的完整路径
        shutil.copy(image_file_path, val_image_dir)  # 将图像文件复制到验证集图像目录

    print("数据集划分完成。")

# json_dir_path = '/home/form_recognition/TabRecSet (CurveTabSet)/TD annotation/chinese'
# image_dir_path = '/home/form_recognition/TabRecSet (CurveTabSet)/image/chinese_all-line'
json_dir_path = '/home/form_recognition/TabRecSet (CurveTabSet)/TD annotation/english'
image_dir_path = '/home/form_recognition/TabRecSet (CurveTabSet)/image/english_all-line'
output_dir_path = '/home/form_recognition/yolo/datasets'
print(f"json目录: {json_dir_path}")
print(f"图片目录: {image_dir_path}")
print(f"输出目录: {output_dir_path}")
split_dataset(json_dir_path, image_dir_path, output_dir_path, train_ratio=0.8)
